import numpy as np
import matplotlib.pyplot as plt
from pyswarm import pso
from scipy.optimize import minimize
from scipy.optimize import Bounds
import time
from tqdm import tqdm

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 定义所有必要的函数和参数
g = 9.8
missile_v = 300
cloud_sink = 3
cloud_r = 10
cloud_t = 20

fake_target = np.array([0, 0, 0])
real_target = np.array([0, 200, 0])
target_r = 7
target_h = 10

M1_start = np.array([20000, 0, 2000])
FY1_start = np.array([17800, 0, 1800])

missile_dir = fake_target - M1_start
missile_dir_norm = np.linalg.norm(missile_dir)
missile_unit_dir = missile_dir / missile_dir_norm
missile_time = missile_dir_norm / missile_v

def missile_position(t):
    return M1_start + missile_unit_dir * missile_v * t

def cloud_position(alpha, v, t_drop, t_burst, tau):
    drop_pos = FY1_start + np.array([v * np.cos(alpha) * t_drop,
                                     v * np.sin(alpha) * t_drop,
                                     0])

    burst_pos = drop_pos + np.array([v * np.cos(alpha) * t_burst,
                                     v * np.sin(alpha) * t_burst,
                                     -0.5 * g * t_burst ** 2])

    return np.array([burst_pos[0],
                     burst_pos[1],
                     max(0, burst_pos[2] - cloud_sink * tau)])

def distance_to_line(point, line_start, line_end):
    line_vec = line_end - line_start
    point_vec = point - line_start
    line_len = np.linalg.norm(line_vec)

    if line_len < 1e-9:
        return np.linalg.norm(point - line_start)

    line_unitvec = line_vec / line_len
    point_vec_scaled = point_vec / line_len
    t = np.dot(line_unitvec, point_vec_scaled)

    if t < 0:
        t = 0
    elif t > 1:
        t = 1

    nearest = line_start + t * line_vec
    return np.linalg.norm(nearest - point)

def sight_blocked(missile_pos, cloud_pos, target_center, target_r, target_h):
    num_points = 8
    target_points = []

    for i in range(num_points):
        angle = 2 * np.pi * i / num_points
        x = target_center[0] + target_r * np.cos(angle)
        y = target_center[1] + target_r * np.sin(angle)
        z = target_center[2]
        target_points.append(np.array([x, y, z]))

    for i in range(num_points):
        angle = 2 * np.pi * i / num_points
        x = target_center[0] + target_r * np.cos(angle)
        y = target_center[1] + target_r * np.sin(angle)
        z = target_center[2] + target_h
        target_points.append(np.array([x, y, z]))

    max_dist = 0
    for target_point in target_points:
        dist = distance_to_line(cloud_pos, missile_pos, target_point)
        if dist > max_dist:
            max_dist = dist
            if max_dist > cloud_r:
                return False
    return max_dist <= cloud_r

def effective_shielding_time(alpha, v, t_drop, t_burst):
    t_bur_total = t_drop + t_burst
    t_cloud_end = t_bur_total + cloud_t
    if t_cloud_end > missile_time:
        t_cloud_end = missile_time
    time_step = 0.5
    effective_time = 0
    for t in np.arange(0, min(cloud_t, missile_time - t_bur_total), time_step):
        t_current = t_bur_total + t
        missile_pos = missile_position(t_current)
        cloud_pos = cloud_position(alpha, v, t_drop, t_burst, t)
        if sight_blocked(missile_pos, cloud_pos, real_target, target_r, target_h):
            effective_time += time_step
    return effective_time

def objective_function(x):
    alpha, v, t_drop, t_burst = x
    if (v < 70 or v > 140 or
            t_drop < 0 or t_burst < 0 or
            t_drop + t_burst > missile_time - cloud_t):
        return 1e6
    shielding_time = effective_shielding_time(alpha, v, t_drop, t_burst)
    if shielding_time == 0:
        missile_pos_at_burst = missile_position(t_drop + t_burst)
        cloud_pos_at_burst = cloud_position(alpha, v, t_drop, t_burst, 0)

        missile_to_target_dist = np.linalg.norm(real_target - missile_pos_at_burst)
        cloud_to_target_dist = np.linalg.norm(real_target - cloud_pos_at_burst)

        penalty = 1000 * (1 / (missile_to_target_dist + cloud_to_target_dist + 1))
        return penalty
    return -shielding_time

lb = [0, 70, 0, 0]
ub = [2 * np.pi, 140, missile_time - cloud_t, cloud_t]

# 定义参数扫描范围
phip_range = np.arange(1.5, 2.05, 0.1)
phig_range = np.arange(2.0, 2.55, 0.1)

# 存储结果的矩阵
fopt_matrix = np.zeros((len(phip_range), len(phig_range)))

# 设置随机种子以确保结果可重现
np.random.seed(42)

print("开始PSO参数扫描...")

# 遍历所有参数组合
for i, phip in enumerate(tqdm(phip_range, desc="扫描phip")):
    for j, phig in enumerate(phig_range):
        try:
            xopt, fopt = pso(objective_function, lb, ub,
                             swarmsize=50, maxiter=100, minstep=1e-4,
                             phip=phip, phig=phig)
            fopt_matrix[i, j] = -fopt
        except Exception as e:
            fopt_matrix[i, j] = np.nan

# 绘制3D曲面图
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

X, Y = np.meshgrid(phig_range, phip_range)
surf = ax.plot_surface(X, Y, fopt_matrix, cmap='plasma', alpha=0.8)

ax.set_xlabel('社会系数 (phig)')
ax.set_ylabel('认知系数 (phip)')
ax.set_zlabel('有效遮蔽时间 (秒)')
ax.set_title('有效遮蔽时间3D曲面图')

fig.colorbar(surf, ax=ax, shrink=0.5, aspect=5)

plt.tight_layout()
plt.savefig('pso_parameter_3d.png', dpi=300, bbox_inches='tight')
plt.show()